"""
投资组合管理工具函数
提供资产配置优化、风险指标计算等功能
"""
import pandas as pd
import numpy as np
from typing import Dict, List, Any, Optional, Union
from loguru import logger
import datetime

from database.data_collector import DataCollector

def optimize_portfolio(portfolio_data: Dict[str, Any]) -> Dict[str, float]:
    """
    优化投资组合配置
    
    参数:
        portfolio_data: 投资组合数据，包含当前持仓、投资策略等信息
        
    返回:
        优化后的资产配置比例
    """
    logger.info("优化投资组合配置")
    
    try:
        # 获取配置数据
        target_stocks = portfolio_data.get("target_stocks", [])
        risk_preference = portfolio_data.get("risk_preference", "balanced")
        current_weights = portfolio_data.get("current_weights", {})
        
        # 如果没有目标股票，则返回空配置
        if not target_stocks:
            logger.warning("没有目标股票，无法优化投资组合")
            return {}
        
        # 获取历史数据计算协方差矩阵
        data_collector = DataCollector()
        start_date = (datetime.datetime.now() - datetime.timedelta(days=365)).strftime("%Y%m%d")
        
        price_data = {}
        for stock_code in target_stocks:
            try:
                stock_data = data_collector.get_daily_data(stock_code, start_date)
                if not stock_data.empty:
                    price_data[stock_code] = stock_data.set_index('日期')['收盘'].astype(float)
            except Exception as e:
                logger.error(f"获取股票 {stock_code} 历史数据失败: {e}")
        
        # 如果获取的价格数据不足，则返回均等权重
        if len(price_data) < len(target_stocks) * 0.8:  # 至少获取到80%的股票数据
            logger.warning("获取的历史价格数据不足，使用均等权重配置")
            equal_weight = 1.0 / len(target_stocks)
            return {stock: equal_weight for stock in target_stocks}
        
        # 创建价格数据框
        prices_df = pd.DataFrame(price_data)
        
        # 计算每日收益率
        returns_df = prices_df.pct_change().dropna()
        
        # 计算协方差矩阵
        cov_matrix = returns_df.cov()
        
        # 计算预期收益率（使用历史平均值作为简单估计）
        expected_returns = returns_df.mean()
        
        # 基于风险偏好设置优化目标
        if risk_preference == "conservative":
            # 最小化风险
            weights = minimize_portfolio_risk(cov_matrix, expected_returns)
        elif risk_preference == "aggressive":
            # 最大化夏普比率
            weights = maximize_sharpe_ratio(cov_matrix, expected_returns)
        else:  # balanced
            # 平衡风险和收益
            weights = optimize_balanced_portfolio(cov_matrix, expected_returns)
        
        # 考虑当前持仓，避免过度调整
        final_weights = adjust_weights_with_current(weights, current_weights)
        
        # 确保权重总和为1
        weight_sum = sum(final_weights.values())
        if weight_sum > 0:
            final_weights = {k: v / weight_sum for k, v in final_weights.items()}
        
        return final_weights
    
    except Exception as e:
        logger.error(f"投资组合优化失败: {e}")
        return {}  # 返回空字典表示优化失败


def minimize_portfolio_risk(cov_matrix: pd.DataFrame, expected_returns: pd.Series) -> Dict[str, float]:
    """
    最小化投资组合风险
    
    参数:
        cov_matrix: 协方差矩阵
        expected_returns: 预期收益率
        
    返回:
        优化后的权重配置
    """
    # 在实际应用中，这里会使用二次规划等算法求解最小方差投资组合
    # 简化版：使用逆波动率加权法
    
    # 计算各股票的波动率
    vol = np.sqrt(np.diag(cov_matrix))
    inv_vol = 1 / vol
    
    # 计算权重
    weights = inv_vol / np.sum(inv_vol)
    
    # 转换为字典
    return {cov_matrix.index[i]: weights[i] for i in range(len(weights))}


def maximize_sharpe_ratio(cov_matrix: pd.DataFrame, expected_returns: pd.Series) -> Dict[str, float]:
    """
    最大化夏普比率
    
    参数:
        cov_matrix: 协方差矩阵
        expected_returns: 预期收益率
        
    返回:
        优化后的权重配置
    """
    # 在实际应用中，这里会使用优化算法求解最大夏普比率投资组合
    # 简化版：使用预期收益率与波动率比值加权
    
    # 计算各股票的波动率
    vol = np.sqrt(np.diag(cov_matrix))
    
    # 计算夏普比率（简化版，假设无风险利率为0）
    sharpe = expected_returns / vol
    
    # 将负夏普比率设为0（不投资预期收益为负的股票）
    sharpe = np.maximum(sharpe, 0)
    
    # 如果所有夏普比率都为0，则使用均等权重
    if np.sum(sharpe) == 0:
        weights = np.ones_like(sharpe) / len(sharpe)
    else:
        # 计算权重
        weights = sharpe / np.sum(sharpe)
    
    # 转换为字典
    return {cov_matrix.index[i]: weights[i] for i in range(len(weights))}


def optimize_balanced_portfolio(cov_matrix: pd.DataFrame, expected_returns: pd.Series) -> Dict[str, float]:
    """
    优化平衡型投资组合
    
    参数:
        cov_matrix: 协方差矩阵
        expected_returns: 预期收益率
        
    返回:
        优化后的权重配置
    """
    # 在实际应用中，这里会使用优化算法求解风险收益平衡的投资组合
    # 简化版：结合最小风险和最大夏普比率两种方法
    
    min_risk_weights = minimize_portfolio_risk(cov_matrix, expected_returns)
    max_sharpe_weights = maximize_sharpe_ratio(cov_matrix, expected_returns)
    
    # 简单加权平均（50%最小风险+50%最大夏普比率）
    balanced_weights = {}
    all_stocks = set(list(min_risk_weights.keys()) + list(max_sharpe_weights.keys()))
    
    for stock in all_stocks:
        min_weight = min_risk_weights.get(stock, 0)
        max_weight = max_sharpe_weights.get(stock, 0)
        balanced_weights[stock] = 0.5 * min_weight + 0.5 * max_weight
    
    # 确保权重总和为1
    weight_sum = sum(balanced_weights.values())
    if weight_sum > 0:
        balanced_weights = {k: v / weight_sum for k, v in balanced_weights.items()}
    
    return balanced_weights


def adjust_weights_with_current(target_weights: Dict[str, float], current_weights: Dict[str, float]) -> Dict[str, float]:
    """
    考虑当前持仓调整目标权重，减少过度调整
    
    参数:
        target_weights: 目标权重配置
        current_weights: 当前权重配置
        
    返回:
        调整后的权重配置
    """
    # 如果当前没有持仓，直接返回目标权重
    if not current_weights:
        return target_weights
    
    adjusted_weights = {}
    
    # 获取所有股票代码
    all_stocks = set(list(target_weights.keys()) + list(current_weights.keys()))
    
    for stock in all_stocks:
        target = target_weights.get(stock, 0)
        current = current_weights.get(stock, 0)
        
        from config.settings import DEFAULT_POSITION_SIZE, MAX_POSITION_PER_STOCK
        # 如果差异小于默认仓位的一半，则保持当前权重
        if abs(target - current) < DEFAULT_POSITION_SIZE / 2:
            adjusted_weights[stock] = current
        else:
            # 否则，向目标权重靠近，但最大调整不超过当前权重的一定比例
            max_adjustment = current * (MAX_POSITION_PER_STOCK / 2)
            adjustment = min(abs(target - current), max_adjustment)
            
            if target > current:
                adjusted_weights[stock] = current + adjustment
            else:
                adjusted_weights[stock] = max(0, current - adjustment)
    
    # 确保权重总和为1
    weight_sum = sum(adjusted_weights.values())
    if weight_sum > 0:
        adjusted_weights = {k: v / weight_sum for k, v in adjusted_weights.items()}
    
    return adjusted_weights


def calculate_risk_metrics(portfolio_data: Dict[str, Any]) -> Dict[str, Any]:
    """
    计算投资组合风险指标
    
    参数:
        portfolio_data: 投资组合数据，包含持仓、历史收益率等信息
        
    返回:
        风险指标计算结果
    """
    logger.info("计算投资组合风险指标")
    
    try:
        # 获取配置数据
        holdings = portfolio_data.get("holdings", {})
        historical_returns = portfolio_data.get("historical_returns", [])
        benchmark_returns = portfolio_data.get("benchmark_returns", [])
        
        # 如果没有持仓或历史收益率数据，则返回基本风险指标
        if not holdings or not historical_returns:
            logger.warning("缺少持仓或历史收益率数据，返回基本风险指标")
            return {
                "volatility": None,
                "max_drawdown": None,
                "sharpe_ratio": None,
                "var_95": None,
                "beta": None,
                "concentration_risk": calculate_concentration_risk(holdings)
            }
        
        # 将历史收益率转换为numpy数组
        returns_array = np.array(historical_returns)
        
        # 计算波动率（年化）
        from config.settings import TRADE_DATES_BUFFER
        volatility = np.std(returns_array) * np.sqrt(TRADE_DATES_BUFFER)
        
        # 计算最大回撤
        cumulative_returns = np.cumprod(1 + returns_array)
        running_max = np.maximum.accumulate(cumulative_returns)
        drawdowns = (cumulative_returns - running_max) / running_max
        max_drawdown = abs(min(drawdowns))
        
        # 计算夏普比率
        from config.settings import TRADE_DATES_BUFFER
        # 假设年化无风险利率为3%，可以在配置中添加
        annual_risk_free_rate = 0.03
        daily_risk_free_rate = annual_risk_free_rate / TRADE_DATES_BUFFER  # 日化无风险利率
        excess_returns = returns_array - daily_risk_free_rate
        sharpe_ratio = np.mean(excess_returns) / np.std(excess_returns) * np.sqrt(TRADE_DATES_BUFFER)
        
        # 计算95%置信区间的VaR
        var_95 = np.percentile(returns_array, 5)
        
        # 计算Beta（如果有基准收益率）
        beta = None
        if benchmark_returns and len(benchmark_returns) == len(historical_returns):
            benchmark_array = np.array(benchmark_returns)
            covariance = np.cov(returns_array, benchmark_array)[0, 1]
            benchmark_variance = np.var(benchmark_array)
            if benchmark_variance > 0:
                beta = covariance / benchmark_variance
        
        # 计算集中度风险
        concentration_risk = calculate_concentration_risk(holdings)
        
        # 其他风险指标
        sortino_ratio = calculate_sortino_ratio(returns_array)
        calmar_ratio = calculate_calmar_ratio(returns_array, max_drawdown)
        
        return {
            "volatility": volatility,
            "max_drawdown": max_drawdown,
            "sharpe_ratio": sharpe_ratio,
            "sortino_ratio": sortino_ratio,
            "calmar_ratio": calmar_ratio,
            "var_95": var_95,
            "beta": beta,
            "concentration_risk": concentration_risk
        }
    
    except Exception as e:
        logger.error(f"风险指标计算失败: {e}")
        return {
            "error": str(e),
            "volatility": None,
            "max_drawdown": None,
            "sharpe_ratio": None,
            "var_95": None,
            "beta": None,
            "concentration_risk": None
        }


def calculate_concentration_risk(holdings: Dict[str, float]) -> float:
    """
    计算持仓集中度风险
    
    参数:
        holdings: 持仓权重
        
    返回:
        集中度风险指标
    """
    # 如果没有持仓，则返回0
    if not holdings:
        return 0.0
    
    # 计算赫芬达尔-赫希曼指数(HHI)
    # HHI是持仓权重平方和，范围为1/N（完全分散）到1（完全集中）
    weights = list(holdings.values())
    hhi = sum([w**2 for w in weights])
    
    return hhi

def calculate_sortino_ratio(returns: np.ndarray, risk_free_rate: Optional[float] = None) -> float:
    """
    计算Sortino比率（只考虑下行风险的夏普比率变体）
    
    
    参数:
        returns: 收益率序列
        risk_free_rate: 无风险利率（默认日化3%）
        
    返回:
        Sortino比率
    """
    # 如果未提供无风险利率，则使用默认值
    if risk_free_rate is None:
        from config.settings import TRADE_DATES_BUFFER
        # 假设年化无风险利率为3%，可以在配置中添加
        annual_risk_free_rate = 0.03
        risk_free_rate = annual_risk_free_rate / TRADE_DATES_BUFFER
    
    excess_returns = returns - risk_free_rate
    
    # 计算下行波动率（只考虑负收益）
    negative_returns = excess_returns[excess_returns < 0]
    
    if len(negative_returns) == 0 or np.std(negative_returns) == 0:
        return np.inf  # 如果没有负收益，则Sortino比率为无穷大
    
    downside_deviation = np.std(negative_returns)
    from config.settings import TRADE_DATES_BUFFER
    sortino_ratio = np.mean(excess_returns) / downside_deviation * np.sqrt(TRADE_DATES_BUFFER)
    
    return sortino_ratio


def calculate_calmar_ratio(returns: np.ndarray, max_drawdown: float) -> float:
    """
    计算Calmar比率（年化收益率与最大回撤的比值）
    
    参数:
        returns: 收益率序列
        max_drawdown: 最大回撤
        
    返回:
        Calmar比率
    """
    if max_drawdown == 0:
        return np.inf  # 如果最大回撤为0，则Calmar比率为无穷大
    
    from config.settings import TRADE_DATES_BUFFER
    annual_return = np.mean(returns) * TRADE_DATES_BUFFER
    calmar_ratio = annual_return / max_drawdown
    
    return float(calmar_ratio)  # 转换为Python float类型


def evaluate_portfolio_performance(portfolio_data: Dict[str, Any]) -> Dict[str, Any]:
    """
    评估投资组合表现
    
    参数:
        portfolio_data: 投资组合数据，包含持仓、历史净值等信息
        
    返回:
        表现评估结果
    """
    logger.info("评估投资组合表现")
    
    try:
        # 获取配置数据
        historical_nav = portfolio_data.get("historical_nav", [])
        benchmark_nav = portfolio_data.get("benchmark_nav", [])
        trades = portfolio_data.get("trades", [])
        
        # 如果没有历史净值数据，则返回基本表现指标
        if not historical_nav:
            logger.warning("缺少历史净值数据，无法评估投资组合表现")
            return {
                "total_return": None,
                "annualized_return": None,
                "win_rate": calculate_win_rate(trades),
                "profit_factor": calculate_profit_factor(trades)
            }
        
        # 计算总收益率
        total_return = historical_nav[-1] / historical_nav[0] - 1
        
        # 计算年化收益率
        days = len(historical_nav)
        from config.settings import TRADE_DATES_BUFFER
        annualized_return = (1 + total_return) ** (TRADE_DATES_BUFFER / days) - 1
        
        # 计算与基准的比较（如果有基准数据）
        alpha = None
        if benchmark_nav and len(benchmark_nav) == len(historical_nav):
            benchmark_return = benchmark_nav[-1] / benchmark_nav[0] - 1
            from config.settings import TRADE_DATES_BUFFER
            benchmark_annualized = (1 + benchmark_return) ** (TRADE_DATES_BUFFER / days) - 1
            alpha = annualized_return - benchmark_annualized
        
        # 计算胜率和盈亏比
        win_rate = calculate_win_rate(trades)
        profit_factor = calculate_profit_factor(trades)
        
        return {
            "total_return": total_return,
            "annualized_return": annualized_return,
            "alpha": alpha,
            "win_rate": win_rate,
            "profit_factor": profit_factor,
            "monthly_returns": calculate_monthly_returns(historical_nav),
            "annual_returns": calculate_annual_returns(historical_nav)
        }
    
    except Exception as e:
        logger.error(f"投资组合表现评估失败: {e}")
        return {
            "error": str(e),
            "total_return": None,
            "annualized_return": None,
            "win_rate": None,
            "profit_factor": None
        }


def calculate_win_rate(trades: List[Dict[str, Any]]) -> Optional[float]:
    """
    计算交易胜率
    
    参数:
        trades: 交易记录列表
        
    返回:
        胜率
    """
    if not trades:
        return None
    
    winning_trades = sum(1 for trade in trades if trade.get("profit", 0) > 0)
    return winning_trades / len(trades)


def calculate_profit_factor(trades: List[Dict[str, Any]]) -> Optional[float]:
    """
    计算盈亏比
    
    参数:
        trades: 交易记录列表
        
    返回:
        盈亏比
    """
    if not trades:
        return None
    
    gross_profit = sum(trade.get("profit", 0) for trade in trades if trade.get("profit", 0) > 0)
    gross_loss = abs(sum(trade.get("profit", 0) for trade in trades if trade.get("profit", 0) < 0))
    
    if gross_loss == 0:
        return float('inf')  # 如果没有亏损交易，则盈亏比为无穷大
    
    return gross_profit / gross_loss


def calculate_monthly_returns(nav_series: List[float]) -> Dict[str, float]:
    """
    计算月度收益率
    
    参数:
        nav_series: 净值序列
        
    返回:
        月度收益率字典
    """
    # 在实际应用中，需要考虑日期信息
    # 这里简化处理，假设nav_series是按照交易日顺序排列的
    
    # 使用配置中的交易日数量
    from config.settings import TRADE_DATES_BUFFER
    # 假设每月交易日为21天，可以在配置中添加
    TRADE_DAYS_PER_MONTH = TRADE_DATES_BUFFER // 12
    
    monthly_returns = {}
    
    for i in range(TRADE_DAYS_PER_MONTH, len(nav_series), TRADE_DAYS_PER_MONTH):
        month_idx = i // TRADE_DAYS_PER_MONTH
        monthly_return = nav_series[i] / nav_series[i-TRADE_DAYS_PER_MONTH] - 1
        monthly_returns[f"Month-{month_idx}"] = monthly_return
    
    return monthly_returns


def calculate_annual_returns(nav_series: List[float]) -> Dict[str, float]:
    """
    计算年度收益率
    
    参数:
        nav_series: 净值序列
        
    返回:
        年度收益率字典
    """
    # 在实际应用中，需要考虑日期信息
    # 这里简化处理，假设nav_series是按照交易日顺序排列的
    
    # 使用配置中的交易日数量
    from config.settings import TRADE_DATES_BUFFER
    
    annual_returns = {}
    
    for i in range(TRADE_DATES_BUFFER, len(nav_series), TRADE_DATES_BUFFER):
        year_idx = i // TRADE_DATES_BUFFER
        annual_return = nav_series[i] / nav_series[i-TRADE_DATES_BUFFER] - 1
        annual_returns[f"Year-{year_idx}"] = annual_return
    
    return annual_returns